import pymysql
import re

DB_CONFIG = {
    'host': 'localhost',
    'user': 'root',
    'password': '1234',
    'database': 'harbin_transport_analysis',
    'port': 3307,
    'charset': 'utf8mb4'
}

def import_dispatch_data(filepath):
    """
    专门用于解析和导入 dispatch_suggestion_analysis.tsv 的函数。
    使用最直接的数据库操作，避免复杂的库依赖。
    """
    print(f"\n--- 正在使用 V2 脚本处理 {filepath} ---")
    
    connection = None
    try:
        # 1. 建立数据库连接
        connection = pymysql.connect(**DB_CONFIG)
        cursor = connection.cursor()
        print("成功连接到 MySQL 数据库。")

        # 2. 创建表（如果不存在）
        table_name = 'dispatch_suggestion_data'
        print(f"准备创建或替换表: {table_name}")
        
        # 为了确保从一个干净的状态开始，总是先删除表
        cursor.execute(f"DROP TABLE IF EXISTS `{table_name}`;")
        
        create_table_query = f"""
        CREATE TABLE `{table_name}` (
            `route_id` VARCHAR(255) NOT NULL,
            `station_id` VARCHAR(255) NOT NULL,
            `hour` INT NOT NULL,
            `onboard` INT,
            `avg_crowd_level` FLOAT,
            PRIMARY KEY (`route_id`, `station_id`, `hour`)
        ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
        """
        cursor.execute(create_table_query)
        print(f"表 '{table_name}' 已成功创建。")

        # 3. 解析文件并插入数据
        print("开始解析文件并插入数据...")
        insert_query = f"""
        INSERT INTO `{table_name}` (route_id, station_id, hour, onboard, avg_crowd_level) 
        VALUES (%s, %s, %s, %s, %s);
        """
        
        records_to_insert = []
        
        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                # 1. 按真正的制表符分割，得到两部分
                parts = line.split('\t', 1)
                if len(parts) == 2:
                    key_part, value_part = parts
                    
                    # 2. 分割 key_part
                    key_splits = key_part.rsplit('_', 1)
                    if len(key_splits) == 2:
                        route_station_part, hour = key_splits
                        
                        route_id, station_id = route_station_part.split('_', 1)

                        # 3. 分割 value_part，它包含 '\\t' 字符串
                        value_splits = value_part.split('\\t')
                        if len(value_splits) == 2:
                            onboard, avg_crowd_level = value_splits
                            records_to_insert.append((route_id, station_id, int(hour), int(onboard), float(avg_crowd_level)))

        if not records_to_insert:
            print("警告：文件中没有找到可解析的数据。")
            return

        # 使用 executemany 进行批量插入，效率更高
        cursor.executemany(insert_query, records_to_insert)
        connection.commit()
        
        print(f"成功插入 {cursor.rowcount} 条记录到 '{table_name}'。")

    except pymysql.MySQLError as e:
        print(f"数据库操作失败: {e}")
        if connection:
            connection.rollback()
    except FileNotFoundError:
        print(f"错误：文件未找到 {filepath}")
    except Exception as e:
        print(f"发生未知错误: {e}")
    finally:
        if connection:
            connection.close()
            print("数据库连接已关闭。")

if __name__ == "__main__":
    # 直接运行此脚本即可导入数据
    import_dispatch_data('mapreduce_jobs/data/dispatch_suggestion_analysis.tsv') 